#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# @Project : chat_model 
# @File    : video.py
# @IDE     : PyCharm 
# @Author  :ZH
# @Time    : 2025/1/4 16:24
from fastapi import Security

from apps.api.v1.video_api import video_router
from apps import logic, model_utils
from apps.api.v1.video_api.schema import VideoModelListResponse, VideoModelListQuery, VideoSessionResponse, \
    VideoSessionCreate, VideoSessionListQuery, VideoDialogueListResponse, VideoDialogueListQuery, DefaultVideoResponse, \
    VideoResponse, VideoQuery
from apps.models import User
from apps.schemas.session_schema import SessionListSchema
from apps.utils import get_current_user
from apps.utils.enumerate import SessionTypeEnum


@video_router.post('/list', name='获取视频模型', response_model=VideoModelListResponse)
async def video_model_list(
        *,
        user: User = Security(get_current_user),
        query: VideoModelListQuery
):
    count, data = await logic.video_model_logic.video_model_list(query=query)
    return {
        'code': 0,
        'message': "成功",
        "count": count,
        "data": data
    }


@video_router.post('/create_session', name='创建会话', response_model=VideoSessionResponse)
async def create_session(
        *,
        user: User = Security(get_current_user),
        query: VideoSessionCreate
):
    if not query.model_id:
        default_model = await logic.video_model_logic.get_default_model()
        if default_model:
            query.model_id = default_model.id
        else:
            return {
                'code': -1,
                "message": "暂无默认可用模型"
            }
    session_info = await logic.session_logic.session_create(
        user=user, session_type=SessionTypeEnum.video, model_id=query.model_id, session_name=query.session_name,
        description=query.description
    )
    return {
        'code': 0,
        'message': "成功",
        'data': session_info
    }


@video_router.post('/session_list', name='获取会话列表', response_model=SessionListSchema)
async def chat_session_list(
        *,
        user: User = Security(get_current_user),
        query: VideoSessionListQuery
):
    count, data = await logic.session_logic.session_list(
        session_type=SessionTypeEnum.video, page=query.page, limit=query.limit, user=user,
        name=query.name
    )
    return {
        'code': 0,
        'message': "成功",
        "count": count,
        "data": data
    }


@video_router.post('/dialogue', name='获取对话历史', response_model=VideoDialogueListResponse)
async def chat_dialogue(
        *,
        user: User = Security(get_current_user),
        query: VideoDialogueListQuery

):
    count, data = await logic.dialogue_logic.dialogue_list(query=query, user=user)
    return {
        'code': 0,
        'message': "成功",
        "count": count,
        "data": data
    }


@video_router.get('/clear_dialogue', name='清空对话历史', response_model=DefaultVideoResponse)
async def clear_dialogue(
        *,
        user: User = Security(get_current_user),
        session_id: int

):
    await logic.dialogue_logic.clear_dialogue(session_id=session_id, user=user)
    return {
        'code': 0,
        'message': "成功"
    }


@video_router.delete('/session', name='删除对话', response_model=DefaultVideoResponse)
async def delete_session(
        *,
        user: User = Security(get_current_user),
        session_id: int
):
    await clear_dialogue(user=user, session_id=session_id)
    await logic.session_logic.session_delete(user=user, session_id=session_id)
    return {
        'code': 0,
        'message': "成功"
    }


@video_router.post('/video', name='视频接口', response_model=VideoResponse)
async def video(
        *,
        user: User = Security(get_current_user),
        query: VideoQuery
):
    if not query.image_url and not query.query:
        return {
            'code': -1,
            "message": "请输入问题/图片地址"
        }

    # 查询session是否存在
    if not query.session_id:
        # 没有session的话，默认自动创建一个session
        session_object = await create_session(user=user, query=VideoSessionCreate())
        if session_object.get('code') == 0:
            query.session_id = session_object.get('data').id
        else:
            return session_object
    session_object = await logic.session_logic.session_object(
        user=user, session_id=query.session_id, session_type=SessionTypeEnum.video
    )
    if not session_object:
        return {
            'code': -1,
            "message": "未找到会话"
        }
    model_id = session_object.model_id
    # 查询模型是否还存在
    model_object = await logic.video_model_logic.get_model_object(model_id=model_id)
    if not model_object:
        return {
            'code': -1,
            "message": "未找到模型"
        }

    # 初始化模型配置
    if model_utils.CONFIG.get(model_object.firm.value) and \
            getattr(model_utils, model_utils.CONFIG.get(model_object.firm.value), None) is not None:
        chat_class = getattr(model_utils, model_utils.CONFIG.get(model_object.firm.value))
        chat_object = chat_class(api_key=model_object.api_key, prompt=session_object.prompt)
        response = await chat_object.video(
            prompt=query.query, base_url=model_object.model_base_url, model_name=model_object.model_name,
            image_url=query.image_url
        )
        # 存储聊天
        await logic.dialogue_logic.create_dialogue(
            user=user, session_id=session_object.id, query=query.query, answer=response, model_id=model_id,
            model_name=model_object.model_name
        )
        return {
            'code': 0,
            'message': "成功",
            'data': {
                'session_id': query.session_id,
                'answer': response
            }
        }

    else:
        return {
            'code': -1,
            "message": "未找到模型配置"
        }


